import pandas as pd
import warnings

import lightgbm as lgb
from sklearn.metrics import accuracy_score
from sklearn.model_selection import train_test_split

warnings.filterwarnings('ignore')

data = pd.read_csv('testa/toUser_train_data.csv')
data = data.fillna(-1)
text_to_value_mapping = {'一加': 2,
                         '华为': 7,
                         'OPPO': 0,
                         'vivo': 1,
                         '小米': 8,
                         '荣耀': 9,
                         '三星': 3,
                         '魅族': 10,
                         '其他': 6,
                         '中兴': 4,
                         '中国移动': 5}

selected_features_brand = data.filter(regex='品牌').columns
# 分别对每一列进行映射转换
for column in selected_features_brand:
    data[column] = data[column].map(text_to_value_mapping)
data = data.fillna(-1)

file_path = 'polynomialFeatures_rules.txt'  # 文件路径

# # 使用with语句打开文件，确保文件在使用后自动关闭
with open(file_path, 'r', encoding='utf-8') as file:
    content = file.read()

lines = content.split("\n")

del_col = []
for line in lines:
    if line.__contains__('^2'):
        line = line.replace('^2', '')
        new_name = line + '^2'
        data[new_name] = data[line].apply(lambda x: x ** 2)
    elif line.__contains__(' '):
        feature_name_list = line.split(' ')
        new_name = feature_name_list[0] + ' ' + feature_name_list[1]
        data[new_name] = data[feature_name_list[0]] * data[feature_name_list[1]]

y = data['新5G终端品牌']
X = data.drop(['新5G终端品牌', '用户标识'], axis=1)

# 划分训练集和验证集
X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.2, random_state=42)

# 将训练集和验证集转换为LightGBM数据集格式
train_data = lgb.Dataset(X_train, label=y_train)
val_data = lgb.Dataset(X_val, label=y_val)
early_stopping_rounds = 30
# 设置LightGBM模型参数
params = {
    'objective': 'multiclass',  # 多分类问题使用'multiclass'
    'metric': 'multi_logloss',  # 多分类问题使用'multi_logloss'作为评估指标
    'num_class': 11,  # 类别数
    'num_leaves': 64,
    'learning_rate': 0.02,
    'min_data_in_leaf': 150,
    'feature_fraction': 0.8,
    'bagging_fraction': 0.7,
    'n_jobs': -1,
    'seed': 1024,
    'early_stopping_rounds': early_stopping_rounds

}

# 训练LightGBM模型，并使用早停法
num_rounds = 1000
model = lgb.train(params, train_data, num_boost_round=num_rounds, valid_sets=val_data)
# 使用训练好的模型进行预测
y_pred = model.predict(X_val)
y_pred = y_pred.argmax(axis=1)  # 对预测结果进行argmax操作，得到类别标签

# 计算准确率
accuracy = accuracy_score(y_val, y_pred)

# 打印准确率
print('Accuracy:', accuracy)

lgb.save_model('lgb.txt', num_iteration=lgb.best_iteration)

input_data = pd.read_csv('input_data.csv')
tm_data = input_data.drop('用户标识', axis=1)
y_pred = model.predict(tm_data, num_iteration=model.best_iteration)
y_pred = y_pred.argmax(axis=1)  # 对预测结果进行argmax操作，得到类别标签
input_data['新5G终端品牌'] = y_pred
res_data = input_data[['用户标识', '新5G终端品牌']]
reversed_dict = {value: key for key, value in text_to_value_mapping.items()}
res_data['新5G终端品牌'] = res_data['新5G终端品牌'].map(reversed_dict)
res_data.to_csv('submit_A_lgb.csv', index=None)
